/* ****************************************************************** **
**    OpenSees - Open System for Earthquake Engineering Simulation    **
**          Pacific Earthquake Engineering Research Center            **
** ****************************************************************** */

//
// 文件: GPUMaterialKernels.cu
// 创建时间: 2025年10月3日
// 说明: GPU材料本构核函数实现
//
// 实现的材料：
//   - Drucker-Prager塑性模型（完整版，支持硬化）
//   - 弹性各向同性（参考）
//   - J2塑性（待实现）
//

// ============================================================
// Performance optimization: Disable kernel debug output
// ============================================================
#ifndef GPU_ENABLE_DEBUG_OUTPUT
#define GPU_ENABLE_DEBUG_OUTPUT 0
#endif

#if GPU_ENABLE_DEBUG_OUTPUT
#define KERNEL_DEBUG_PRINT(...) KERNEL_DEBUG_PRINT(__VA_ARGS__)
#else
#define KERNEL_DEBUG_PRINT(...) ((void)0)
#endif

#include "GPUMaterialKernels.h"
#include <math.h>

//==============================================================================
// 辅助设备函数实现
//==============================================================================

/**
 * @brief 计算应力不变量
 */
__device__ void computeStressInvariants(
    const double* stress,
    double& I1,
    double& J2)
{
    // 应力分量：stress = [σxx, σyy, σzz, σxy, σyz, σzx]
    double sxx = stress[0];
    double syy = stress[1];
    double szz = stress[2];
    double sxy = stress[3];
    double syz = stress[4];
    double szx = stress[5];

    // 第一不变量：I1 = tr(σ) = σxx + σyy + σzz
    I1 = sxx + syy + szz;

    // 平均应力
    double p = I1 / 3.0;

    // 偏应力
    double s[6];
    s[0] = sxx - p;
    s[1] = syy - p;
    s[2] = szz - p;
    s[3] = sxy;
    s[4] = syz;
    s[5] = szx;

    // 第二偏应力不变量：J2 = 1/2 * s:s
    J2 = 0.5 * (s[0]*s[0] + s[1]*s[1] + s[2]*s[2]) +
         s[3]*s[3] + s[4]*s[4] + s[5]*s[5];
}

/**
 * @brief 计算弹性刚度矩阵（3D各向同性）
 */
__device__ void computeElasticD(
    double K,  // 体积模量
    double G,  // 剪切模量
    double* D  // 输出：6x6刚度矩阵（按行存储）
)
{
    // 清零
    for (int i = 0; i < 36; i++) {
        D[i] = 0.0;
    }

    // Lamé常数
    double lambda = K - 2.0 * G / 3.0;

    // 对角块（正应力-正应力）
    D[0*6 + 0] = lambda + 2.0 * G;  // D11
    D[0*6 + 1] = lambda;             // D12
    D[0*6 + 2] = lambda;             // D13

    D[1*6 + 0] = lambda;             // D21
    D[1*6 + 1] = lambda + 2.0 * G;  // D22
    D[1*6 + 2] = lambda;             // D23

    D[2*6 + 0] = lambda;             // D31
    D[2*6 + 1] = lambda;             // D32
    D[2*6 + 2] = lambda + 2.0 * G;  // D33

    // 剪应力项
    D[3*6 + 3] = G;  // D44 (σxy)
    D[4*6 + 4] = G;  // D55 (σyz)
    D[5*6 + 5] = G;  // D66 (σzx)
}

//==============================================================================
// Drucker-Prager材料本构实现
//==============================================================================

/**
 * @brief Drucker-Prager塑性材料应力更新（Return Mapping算法）
 *
 * 算法流程：
 * 1. 弹性预测（Elastic Predictor）
 * 2. 检查屈服（Yield Check）
 * 3. 塑性修正（Plastic Corrector - Return Mapping）
 * 4. 更新状态变量
 * 5. 计算切线刚度（Consistent Tangent）
 */
__device__ void computeDruckerPragerStress(
    const double* strain_increment,     // 应变增量 Δε
    const DruckerPragerParams& params,  // 材料参数
    MaterialState& state,               // 材料状态（输入输出）
    double* stress,                     // 输出应力
    double* D_tangent                   // 输出切线刚度
)
{
    const double TOL = 1.0e-10;
    const int MAX_ITER = 50;

    // ========== Step 1: 弹性预测 ==========

    // 计算弹性刚度矩阵
    double D_elastic[36];
    computeElasticD(params.K, params.G, D_elastic);

    // 弹性应力增量：Δσ = D : Δε
    double stress_trial[6] = {0, 0, 0, 0, 0, 0};
    for (int i = 0; i < 6; i++) {
        stress_trial[i] = state.stress[i];  // 从上一步应力开始
        for (int j = 0; j < 6; j++) {
            stress_trial[i] += D_elastic[i*6 + j] * strain_increment[j];
        }
    }

    // 计算trial stress的不变量
    double I1_trial, J2_trial;
    computeStressInvariants(stress_trial, I1_trial, J2_trial);

    // ========== Step 2: 检查屈服 ==========

    // 当前硬化参数
    double alpha = state.alpha;  // 累积塑性应变

    // 硬化函数：k = k0 + (kinf - k0) * (1 - exp(-delta * alpha))
    double k = params.sigma_y;
    if (params.delta1 > 0) {
        k = params.Ko + (params.Kinf - params.Ko) *
            (1.0 - exp(-params.delta1 * alpha));
    }

    // 屈服函数：f = sqrt(J2) + rho * I1 / 3 - k
    double sqrt_J2 = sqrt(fmax(J2_trial, 0.0));
    double f_trial = sqrt_J2 + params.rho * I1_trial / 3.0 - k;

    // ========== 弹性状态：直接返回 ==========
    if (f_trial <= TOL) {
        // 更新应力
        for (int i = 0; i < 6; i++) {
            stress[i] = stress_trial[i];
            state.stress[i] = stress_trial[i];
        }

        // 切线刚度 = 弹性刚度
        for (int i = 0; i < 36; i++) {
            D_tangent[i] = D_elastic[i];
        }
        return;
    }

    // ========== Step 3: 塑性修正（Return Mapping）==========

    // 初始化塑性乘子
    double delta_lambda = 0.0;

    // Newton-Raphson迭代求解塑性乘子
    for (int iter = 0; iter < MAX_ITER; iter++) {
        // 更新硬化参数
        double alpha_new = alpha + delta_lambda;
        double k_new = params.Ko + (params.Kinf - params.Ko) *
                       (1.0 - exp(-params.delta1 * alpha_new));
        double dk_dalpha = params.delta1 * (params.Kinf - params.Ko) *
                          exp(-params.delta1 * alpha_new);

        // Return mapping修正后的应力
        double p_trial = I1_trial / 3.0;
        double p = p_trial - params.K * params.rho * delta_lambda;
        double q = sqrt_J2 - 3.0 * params.G * delta_lambda;

        if (q < 0) q = 0;  // 避免负值

        // 残差：f = q + rho * p - k_new
        double f = q + params.rho * p - k_new;

        // 收敛检查
        if (fabs(f) < TOL) {
            break;
        }

        // 导数：df/d(delta_lambda)
        double df_dlambda = -3.0 * params.G - params.rho * params.K * params.rho - dk_dalpha;

        // Newton更新
        delta_lambda -= f / df_dlambda;

        // 确保delta_lambda非负
        if (delta_lambda < 0) {
            delta_lambda = 0;
        }
    }

    // ========== Step 4: 更新应力和状态变量 ==========

    // 修正后的应力不变量
    double p_trial = I1_trial / 3.0;
    double p_corrected = p_trial - params.K * params.rho * delta_lambda;
    double q_corrected = sqrt_J2 - 3.0 * params.G * delta_lambda;

    if (q_corrected < 0) q_corrected = 0;

    // 重构应力分量
    // p = (σxx + σyy + σzz) / 3
    // q = sqrt(3/2 * s:s) = sqrt(J2)

    double ratio = (q_corrected > TOL) ? q_corrected / sqrt_J2 : 0.0;

    for (int i = 0; i < 6; i++) {
        double s_trial_i = stress_trial[i] - (i < 3 ? I1_trial / 3.0 : 0.0);
        double s_corrected_i = ratio * s_trial_i;

        if (i < 3) {
            stress[i] = p_corrected + s_corrected_i;
        } else {
            stress[i] = s_corrected_i;
        }

        state.stress[i] = stress[i];
    }

    // 更新累积塑性应变
    state.alpha += delta_lambda;

    // ========== Step 5: 计算一致切线刚度 ==========
    // 简化版本：使用弹塑性刚度
    // 完整版本需要计算 D_ep = D_e - (D_e : df/dsigma) x (dg/dsigma : D_e) / H

    // 这里使用简化版本：连续切线刚度
    if (delta_lambda > TOL) {
        // 塑性状态：使用降低的刚度
        double beta = 1.0 / (1.0 + 3.0 * params.G * delta_lambda / (sqrt_J2 + TOL));
        for (int i = 0; i < 36; i++) {
            D_tangent[i] = beta * D_elastic[i];
        }
    } else {
        // 弹性状态
        for (int i = 0; i < 36; i++) {
            D_tangent[i] = D_elastic[i];
        }
    }
}

//==============================================================================
// 弹性各向同性材料（参考实现）
//==============================================================================

__device__ void computeElasticIsotropicStress(
    const double* strain,
    const ElasticIsotropicParams& params,
    double* stress,
    double* D_tangent
)
{
    // 计算体积模量和剪切模量
    double K = params.E / (3.0 * (1.0 - 2.0 * params.nu));
    double G = params.E / (2.0 * (1.0 + params.nu));

    // 计算刚度矩阵
    computeElasticD(K, G, D_tangent);

    // 计算应力：σ = D : ε
    for (int i = 0; i < 6; i++) {
        stress[i] = 0.0;
        for (int j = 0; j < 6; j++) {
            stress[i] += D_tangent[i*6 + j] * strain[j];
        }
    }
}

//==============================================================================
// 材料调度器
//==============================================================================

__device__ void computeMaterialStress(
    GPUMaterialType matType,
    const double* strain_increment,
    const void* params,
    MaterialState& state,
    double* stress,
    double* D_tangent
)
{
    switch (matType) {
        case GPUMaterialType::DRUCKER_PRAGER:
            computeDruckerPragerStress(
                strain_increment,
                *static_cast<const DruckerPragerParams*>(params),
                state,
                stress,
                D_tangent
            );
            break;

        case GPUMaterialType::ELASTIC_ISOTROPIC:
            computeElasticIsotropicStress(
                strain_increment,
                *static_cast<const ElasticIsotropicParams*>(params),
                stress,
                D_tangent
            );
            break;

        default:
            // 不支持的材料类型
            for (int i = 0; i < 6; i++) stress[i] = 0.0;
            for (int i = 0; i < 36; i++) D_tangent[i] = 0.0;
            break;
    }
}
